package beast.evolution.tree.coalescent;
import beast.core.Description;
import beast.core.Input;
import beast.core.parameter.RealParameter;

import java.util.Collections;
import java.util.List;

/**
 * This population function supplies effective population size / rate of coalescence for a SIR demographic model
 * SIR trajectories are simulated deterministically
 * Transmission is frequency-dependent
 *
 * For CoalescentLikelihood I need to implement:
 * getIntegral(startTime, finishTime)
 * getPopSize(finishTime)
 * getIntensity(time)
 *
 * @author Trevor Bedford
 * @author Alexei Drummond
 * @author Denise Kuehnert
 */
@Description("Coalescent intervals for an SIR population.")
public class SIRPopulation extends PopulationFunction.Abstract {

    public Input<RealParameter> transmissionParameter = new Input<RealParameter>("transmissionRate", "Rate of contact between hosts: beta");
    public Input<RealParameter> recoveryParameter = new Input<RealParameter>("recoveryRate", "Rate of recovery of a host: gamma");
    public Input<RealParameter> susceptibleParameter = new Input<RealParameter>("presentDayS", "Susceptible population at present");
    public Input<RealParameter> infectedParameter = new Input<RealParameter>("presentDayI", "Infected population at present");
    public Input<RealParameter> recoveredParameter = new Input<RealParameter>("presentDayR", "Recovery population at present");

    // Public stuff
    public void initAndValidate() throws Exception {
        if (transmissionParameter.get() != null) {
            transmissionParameter.get().setBounds(0.0, Double.POSITIVE_INFINITY);
        }
        if (recoveryParameter.get() != null) {
            recoveryParameter.get().setBounds(0.0, Double.POSITIVE_INFINITY);
        }
        if (susceptibleParameter.get() != null) {
            susceptibleParameter.get().setBounds(0.0, Double.POSITIVE_INFINITY);
        }
        if (infectedParameter.get() != null) {
            infectedParameter.get().setBounds(0.0, Double.POSITIVE_INFINITY);
        }
        if (recoveredParameter.get() != null) {
            recoveredParameter.get().setBounds(0.0, Double.POSITIVE_INFINITY);
        }
    }

    /**
    * @return transmission rate
    */
    public double getTransmission() {
        return transmissionParameter.get().getValue();
    }

    /**
    * @return recovery rate
    */
    public double getRecovery() {
        return recoveryParameter.get().getValue();
    }

    /**
    * @return susceptible population at present
    */
    public double getSusceptible() {
        return susceptibleParameter.get().getValue();
    }

    /**
    * @return susceptible population at present
    */
    public double getInfected() {
        return infectedParameter.get().getValue();
    }

    /**
    * @return susceptible population at present
    */
    public double getRecovered() {
        return recoveredParameter.get().getValue();
    }

    // Implementation of abstract methods

    // return population size, this is the inverse of the coalescence rate
    // todo pull from dynamical timeseries
    public double getPopSize(double t) {
        return 1.0;
    }

    // return coalescence rate
    // todo pull from dynamical timeseries
    public double getIntensity(double t) {
        return 1.0;
    }

    // return time matching to a particular population size
    // todo pull from dynamical timeseries
    public double getInverseIntensity(double x) {
        return 1.0;
    }

    // return coalescence integral of coalescence rates
    // todo pull from dynamical timeseries
    public double getIntegral(double start, double finish) {
        return 1.0;
    }

    // Implementation of abstract methods

    public List<String> getParameterIds() {
        return Collections.singletonList(transmissionParameter.get().getID());
    }

}
